import torch
from torch import nn


class Dense(nn.Module):
    def __init__(self, dims):
        super(Dense, self).__init__()
        self.input_size, self.hidden_size, self.output_size = dims
        self.fc1 = nn.Linear(self.input_size, self.hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(self.hidden_size, self.output_size)

    def forward(self, x):
        # 输入x.shape = (B, T, N, C)
        return self.fc2(self.relu(self.fc1(x)))
